package main

import (
	"code.google.com/p/gochess/movement"
	"code.google.com/p/gochess/notation/fen"
	"code.google.com/p/gochess/notation/lan"
	pos "code.google.com/p/gochess/position"
	"code.google.com/p/gochess/search"
	"flag"
	"fmt"
	humanize "github.com/dustin/go-humanize"
	"os"
	"runtime/pprof"
)

const maxDepth = 31

func showUsage() {
	fmt.Fprintf(os.Stderr, "usage: perft [-d n] [-fen notation] [-g] [-m move] [-h]\n")
	flag.PrintDefaults()
}

func main() {

	cpuprofile := flag.String("cpuprofile", "", "write cpu profile to file")
	depth := flag.Int("d", 0, "search depth")
	fenPosition := flag.String("fen", fen.InitialPosition, "FEN notation for position to search")
	usage := flag.Bool("h", false, "show usage information")
	useGochess := flag.Bool("g", false, "use gochess engine instead of dedicated perft routine")
	flgMove := flag.String("m", "", "last move played in lan")

	flag.Parse()

	if *usage {
		showUsage()
		return
	}
	if *cpuprofile != "" {
		f, err := os.Create(*cpuprofile)
		if err != nil {
			fmt.Printf("Could not create file %s.\n", *cpuprofile)
			os.Exit(1)
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	p, err := fen.Parse(*fenPosition)
	if err != nil {
		fmt.Fprintf(os.Stderr, "invalid fen: %s\n", err)
		os.Exit(1)
	}

	if *depth >= maxDepth && *depth < 1 {
		fmt.Fprintf(os.Stderr, "Invalid depth, must be between 1 and %d\n", maxDepth-1)
		os.Exit(1)
	}
	var lastMove pos.Move
	if *flgMove != "" {
		if lastMove, err = lan.ParseMove(p, *flgMove); err != nil {
			fmt.Fprintf(os.Stderr, "invalid last move: %v\n", err)
			os.Exit(1)
		}
	}
	printPerftResults(p, int32(*depth), lastMove, *useGochess)
}

type nodeCount struct {
	move  pos.Move
	nodes int64
	fen   string
}

func count(p *pos.Position, moves []pos.Move, pf func(pos.Move) int64, out chan<- nodeCount) {

	for _, move := range moves {
		uMove := pos.MakeMove(p, move)
		nodes := pf(move)
		out <- nodeCount{move, nodes, fen.Sprint(p)}
		pos.UnmakeMove(p, uMove)
	}
	close(out)
}

var (
	moveCache      = [maxDepth]movement.Movelist{}
	legalMoveCache = [maxDepth]movement.Movelist{}
)

const tableFormat = "%6v %11v  %v\n"

func printPerftResults(p *pos.Position, depth int32, lastMove pos.Move, useGochess bool) {

	var moves []pos.Move
	var pf func(pos.Move) int64
	if useGochess {
		moves = search.RootMoves(p, lastMove)
		pf = func(move pos.Move) int64 {
			return search.Perft(p, move, depth-1)
		}
	} else {
		moves = movement.GenerateLegalMoves(p)
		pf = func(move pos.Move) int64 {
			return perft(p, depth-1)
		}
	}
	out := make(chan nodeCount)
	go count(p, moves, pf, out)
	var sum int64
	fmt.Printf(tableFormat, "move", "nodes", "position after move")
	for node := range out {
		fmt.Printf(tableFormat, lan.SprintMove(p, node.move),
			humanize.Comma(node.nodes), node.fen)
		sum += node.nodes
	}
	fmt.Printf(tableFormat, "Total", humanize.Comma(sum), "")
}

func perft(p *pos.Position, depth int32) int64 {

	if depth <= 0 {
		return 1
	}
	var nodes int64
	var uMove pos.MoveUnmaking
	moves := &(moveCache[depth])
	moves.Clear()
	movement.GenerateMoves(p, moves)
	stm := p.SideToMove()
	for move := moves.Next(); move != pos.NoMove; move = moves.Next() {
		uMove = pos.MakeMove(p, move)
		if !movement.IsCheck(p, stm) {
			nodes += perft(p, depth-1)
		}
		pos.UnmakeMove(p, uMove)
	}
	return nodes
}
